{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wjRvJ8zBftL3"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.CustomCallback at 0x7f9b51e87c40>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from stable_baselines3.common.callbacks import BaseCallback\n",
    "\n",
    "\n",
    "#Callback语法\n",
    "class CustomCallback(BaseCallback):\n",
    "\n",
    "    def __init__(self, verbose=0):\n",
    "        super().__init__(verbose)\n",
    "\n",
    "        #可以访问的变量\n",
    "        #self.model\n",
    "        #self.training_env\n",
    "        #self.n_calls\n",
    "        #self.num_timesteps\n",
    "        #self.locals\n",
    "        #self.globals\n",
    "        #self.logger\n",
    "        #self.parent\n",
    "\n",
    "    def _on_training_start(self) -> None:\n",
    "        #第一个rollout开始前调用\n",
    "        pass\n",
    "\n",
    "    def _on_rollout_start(self) -> None:\n",
    "        #rollout开始前\n",
    "        pass\n",
    "\n",
    "    def _on_step(self) -> bool:\n",
    "        #env.step()之后调用,返回False后停止训练\n",
    "        return True\n",
    "\n",
    "    def _on_rollout_end(self) -> None:\n",
    "        #更新参数前调用\n",
    "        pass\n",
    "\n",
    "    def _on_training_end(self) -> None:\n",
    "        #训练结束前调用\n",
    "        pass\n",
    "\n",
    "\n",
    "CustomCallback()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7ILY0AkFfzPJ"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080\">   0%</span> <span style=\"color: #3a3a3a; text-decoration-color: #3a3a3a\">━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━</span> <span style=\"color: #008000; text-decoration-color: #008000\">0/8,000 </span> [ <span style=\"color: #808000; text-decoration-color: #808000\">0:00:00</span> &lt; <span style=\"color: #008080; text-decoration-color: #008080\">-:--:--</span> , <span style=\"color: #800000; text-decoration-color: #800000\">? it/s</span> ]\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[35m   0%\u001b[0m \u001b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0/8,000 \u001b[0m [ \u001b[33m0:00:00\u001b[0m < \u001b[36m-:--:--\u001b[0m , \u001b[31m? it/s\u001b[0m ]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">20\n",
       "</pre>\n"
      ],
      "text/plain": [
       "20\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">40\n",
       "</pre>\n"
      ],
      "text/plain": [
       "40\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">60\n",
       "</pre>\n"
      ],
      "text/plain": [
       "60\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">80\n",
       "</pre>\n"
      ],
      "text/plain": [
       "80\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">100\n",
       "</pre>\n"
      ],
      "text/plain": [
       "100\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<stable_baselines3.ppo.ppo.PPO at 0x7f9b51eb7310>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from stable_baselines3 import PPO\n",
    "\n",
    "\n",
    "#让训练只执行N步的callback\n",
    "class LimitStepCallback(BaseCallback):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__(verbose=0)\n",
    "        self.call_count = 0\n",
    "\n",
    "    def _on_step(self):\n",
    "        self.call_count += 1\n",
    "\n",
    "        if self.call_count % 20 == 0:\n",
    "            print(self.call_count)\n",
    "\n",
    "        if self.call_count >= 100:\n",
    "            return False\n",
    "\n",
    "        return True\n",
    "\n",
    "\n",
    "model = PPO('MlpPolicy', 'CartPole-v1', verbose=0)\n",
    "\n",
    "model.learn(8000, callback=[LimitStepCallback()], progress_bar=True)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "4_callbacks_hyperparameter_tuning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python [conda env:cuda117]",
   "language": "python",
   "name": "conda-env-cuda117-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
